import hashlib, os
from support.utils.toolkit import Log
from enum import Enum
from support.dbms.helper import DBMS
from support.ai.chroma import VectorDB 
from support.files.text import TextDoc
class FILE_STATUS(Enum):
    NOTEMBED = 0
    EMBEDDED = 1
    
class EMBED:
    TAG = "EMBED"
    DEFAULT_LENGTH = 4096
    DEFAULT_EMBEDDINGS = "EntropyYue/jina-embeddings-v2-base-zh:latest"
    DATABASE_NAME = "embed"
    TABLE_NAME_BOOKS = "books"
    COLLECTION_NAME="onebook"
    RESOURCE_PATH="/books"

    __instance__ = None
    def __new__(cls, *args, **kwargs):
        if(not cls.__instance__): cls.__instance__ = object.__new__(cls)
        return cls.__instance__

    def getInstance():
        return EMBED.__instance__

    def __init__(self):
        self.helper = VectorDB(EMBED.COLLECTION_NAME, embeddings=EMBED.DEFAULT_EMBEDDINGS)


    def getTables(self):
        return [{
            "name":EMBED.TABLE_NAME_BOOKS, 
            "keys":[
                {"label":"NAME", "type":"TEXT"},
                {"label":"PATH", "type":"TEXT"},
                {"label":"HASH", "type":"CHAR", "length":32},
                {"label":"EMBEDED", "type":"INT"}
             ]
         }]

    def __calcFileMd5Value__(self, filepath):
        hashMd5 = hashlib.md5()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(EMBED.DEFAULT_LENGTH), b""):
                hashMd5.update(chunk)
        return hashMd5.hexdigest()

    def __isNewFile__(self, hashvalue, oper):
        fileid = None
        status = None
        isEmbedding = None
        ret = DBMS.query(oper, EMBED.TABLE_NAME_BOOKS, {"ID":None, "EMBEDED":None, "HASH":hashvalue})
        if(ret):
            fileid= ret[0]["ID"]
            isEmbeded = ret[0]["EMBEDED"]
        if (fileid is None or isEmbeded != FILE_STATUS.EMBEDDED.value):
            status = FILE_STATUS.NOTEMBED
        elif(fileid):
            status = FILE_STATUS.EMBEDDED
        return (fileid, status)


    def exec(self, oper):
        for root, dirs, files in os.walk(EMBED.RESOURCE_PATH):
            for file in files:
                filepath = os.path.join(root, file)
                hashvalue = self.__calcFileMd5Value__(filepath)
                fileStatus = self.__isNewFile__(hashvalue, oper)
                if(fileStatus[1] != FILE_STATUS.EMBEDDED): 
                    doc = TextDoc().getSentenceBoundaryFrom(filepath)
                    metadata = {"source": file}
                    try:
                        while(True):
                            seg = next(doc)
                            if(seg is None): break
                            for text in seg: self.helper.addData(metadata, text) 
                    except StopIteration:
                        DBMS.insert(oper, EMBED.TABLE_NAME_BOOKS, {"NAME":file, "PATH":filepath, "HASH":hashvalue, "EMBEDED":FILE_STATUS.EMBEDDED.value})
                    except Exception as e:
                        Log.e(str(e))
